Rewrite Gemma4 scannable block by aireenmei · Pull Request #4530 · AI-Hypercomputer/maxtext

aireenmei · 2026-07-18T06:17:18Z

Description

Currently theimplementation of Gemma4ScannableBlock makes XLA treat the whole 6-layer block as a "big layer". The all-gathers for all 6 layers are run at the beginning of each block. Similarly, in backward pass, remat runs the whole block. This causes elevated HBM usage and more exposed collectives.

After this rewrite:

5 local layers via jax.lax.scan, each per-layer jax.checkpoint
length-1 scan for global layer, also wrapped in jax.checkpoint
remove block-level remat and unroll outer block scan

Tests

added unit test
Before rewrite: xprof for pdbs=4. xprof for pdbs=8
After rewrite: xprof for pdbs=4. xprof for pdbs=8 + vocab_tiling=8 (OOM without vocab_tiling).
Before rewrite, the HBM peak is not on vocab peak(screenshot). After rewrite, vocab peak becomes the highest (screenshot) due to other peaks become smaller, also by having remat boundary at layer level adds HBM usage and overlap with the vocab peak. Therefore, without vocab tiling, the peak HBM usage doesn't change much compared with original. But when combined with vocab_tiling, we can effectively lower the HBM peak a lot (70 GB -> 40 GB for pdbs=4).

Checklist

Before submitting this PR, please make sure (put X in square brackets):

I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
I have necessary comments in my code, particularly in hard-to-understand areas.
I have run end-to-end tests tests and provided workload links above if applicable.
I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

github-actions · 2026-07-18T09:22:12Z

🤖 CI Failure Investigation Report

I have analyzed the recent test failures in the CI pipeline and identified the following:

🔍 What Failed

Job/Matrix: code_quality_check / Static code-quality checkers
Failing Test: pre-commit checks on just the files that have changed
Error: Pylint and pyink code style violations

🪵 Error Details & Stack Trace

# Missing function docstring in Gemma4ScannableBlock helper method
src/maxtext/models/gemma4.py:496:2: C0116: Missing function or method docstring (missing-function-docstring)

# Line too long violations in model_test.py
tests/unit/model_test.py:238:0: C0301: Line too long (131/125) (line-too-long)
tests/unit/model_test.py:283:0: C0301: Line too long (131/125) (line-too-long)

# Formatting non-compliance reported by pyink
would reformat tests/unit/model_test.py
would reformat src/maxtext/models/gemma4.py
would reformat tests/unit/nnx_scan_test.py
would reformat tests/unit/nnx_decoders_test.py

💡 Root Cause Analysis & Context

Confidence: high (confirmed cause)

The failure is caused by code style and quality violations introduced in this PR. Specifically:

In src/maxtext/models/gemma4.py, the new helper method _apply_local_layers was added without a docstring, triggering the pylint missing-function-docstring check.
In tests/unit/model_test.py, the lines asserting logits shapes exceed the maximum length of 125 characters configured for pylint (each was 131 characters).
The files tests/unit/model_test.py, src/maxtext/models/gemma4.py, tests/unit/nnx_scan_test.py, and tests/unit/nnx_decoders_test.py had spacing and line wraps that did not comply with the pyink formatting rules (indentation 2, line length 122).

These issues are purely style/quality regressions and do not affect behavioral correctness. They have been fully corrected and verified locally using pre-commit run.

🛠️ Recommended Fix

Apply the formatting changes and add the missing docstring to src/maxtext/models/gemma4.py. The required diffs are:

diff --git a/src/maxtext/models/gemma4.py b/src/maxtext/models/gemma4.py
index b974ec5..f626b8e 100644
--- a/src/maxtext/models/gemma4.py
+++ b/src/maxtext/models/gemma4.py
@@ -453,9 +453,7 @@ class Gemma4ScannableBlock(nnx.Module):
 
     pattern_length = len(GEMMA4_ATTENTION_PATTERN)
     if not 0 <= num_of_layers <= pattern_length:
-      raise ValueError(
-          f"Gemma4ScannableBlock must contain between 0 and {pattern_length} layers; got {num_of_layers}."
-      )
+      raise ValueError(f"Gemma4ScannableBlock must contain between 0 and {pattern_length} layers; got {num_of_layers}.")
 
     # Pattern is 5 local, 1 global.
     self.num_local = min(5, num_of_layers)
@@ -505,6 +503,8 @@ class Gemma4ScannableBlock(nnx.Module):
       bidirectional_mask=None,
       attention_metadata=None,
   ):
+    """Applies local attention layers sequentially using scan."""
+
     def apply_layer(layer, carry):
       layer_out = layer(
           carry,
@@ -654,9 +652,7 @@ class Gemma4ScannableBlock(nnx.Module):
         offload_names = maxtext_utils.get_save_and_offload_names(cfg)
         if offload_names[0] or offload_names[1]:
           save_names, offload_to_device = offload_names
-          global_remat_policy = jax.checkpoint_policies.save_only_these_names(
-              *(save_names + offload_to_device)
-          )
+          global_remat_policy = jax.checkpoint_policies.save_only_these_names(*(save_names + offload_to_device))
 
         if self.apply_internal_remat and self.config.remat_policy != "none":
           prevent_cse = maxtext_utils.should_prevent_cse_in_remat(self.config)
@@ -669,9 +667,7 @@ class Gemma4ScannableBlock(nnx.Module):
         # Carry state through the loop instead of returning a stacked [1, ...]
         # scan result: slicing that result previously introduced a bitcast
         # between device and pinned-host memory under offload remat.
-        with xla_metadata.set_xla_metadata(
-            **{"skip-simplify-while-loops_trip-count-one": "true"}
-        ):
+        with xla_metadata.set_xla_metadata(**{"skip-simplify-while-loops_trip-count-one": "true"}):
           (y, global_state), _ = jax.lax.scan(
               scan_global_layer,
               (y, state_g),
diff --git a/tests/unit/model_test.py b/tests/unit/model_test.py
index 7c3d4ea..dd679ff 100644
--- a/tests/unit/model_test.py
+++ b/tests/unit/model_test.py
@@ -235,7 +235,9 @@ class TestModel(unittest.TestCase):
         enable_dropout=False,
         model_mode=MODEL_MODE_TRAIN,
     )
-    self.assertEqual(logits.shape, (new_config.global_batch_size_to_train_on, new_config.max_target_length, new_config.vocab_size))
+    self.assertEqual(
+        logits.shape, (new_config.global_batch_size_to_train_on, new_config.max_target_length, new_config.vocab_size)
+    )
 
   def test_gemma4_model_linen(self):
     """Test the shared Gemma4 scannable block on the linen (ToLinen) path.
@@ -280,7 +282,9 @@ class TestModel(unittest.TestCase):
         rngs={"aqt": self.rng},
     )
     logits = logits[0] if isinstance(logits, tuple) else logits
-    self.assertEqual(logits.shape, (new_config.global_batch_size_to_train_on, new_config.max_target_length, new_config.vocab_size))
+    self.assertEqual(
+        logits.shape, (new_config.global_batch_size_to_train_on, new_config.max_target_length, new_config.vocab_size)
+    )

codecov · 2026-07-21T06:44:26Z

Codecov Report

❌ Patch coverage is 79.79798% with 20 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/models/gemma4.py	80.00%	6 Missing and 8 partials ⚠️
src/maxtext/layers/nnx_scan.py	80.00%	1 Missing and 3 partials ⚠️
src/maxtext/layers/nnx_decoders.py	77.77%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

aireenmei · 2026-07-21T21:13:35Z

I add pull ready to generate internal cl for testing. Won't submit until PR approved

github-actions · 2026-07-22T00:17:48Z

🤖 Hi @RissyRan, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

github-actions · 2026-07-22T00:20:57Z

🤖 I'm sorry @RissyRan, but I was unable to process your request. Please see the logs for more details.

RissyRan

Thanks for the change! Overall it looks good. I am a little bit lost in gemma4.py if/else branches, wondering if we could simplify a little bit?

RissyRan · 2026-07-23T00:02:15Z

+    if not 0 <= num_of_layers <= pattern_length:
+      raise ValueError(f"Gemma4ScannableBlock must contain between 0 and {pattern_length} layers; got {num_of_layers}.")
+
+    # Pattern is 5 local, 1 global.


It seems the pattern has been defined. Will something like this work?

active_pattern = GEMMA4_ATTENTION_PATTERN[:num_of_layers] self.num_local = sum(1 for t in active_pattern if t == AttentionType.LOCAL_SLIDING) self.num_global = sum(1 for t in active_pattern if t == AttentionType.GLOBAL)

RissyRan · 2026-07-23T00:06:53Z

+          metadata_axis_name="local_layers",
+          rngs=self.rngs,
+      )
+    else:


Will this be triggered for Gemma4?

RissyRan · 2026-07-23T00:09:01Z

-      current_kv = kv_cache[layer_id] if kv_cache is not None else None
-      y, new_kv = getattr(self, f"layers_{layer_id}")(
+
+    if kv_cache is not None:


I guess you meet issue when calling kv_cache[layer_id]?

Wondering if you tested for decoding with the change end-to-end and see reasonable outputs?

RissyRan · 2026-07-23T00:19:23Z

+          stacked_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, scan_axis), stacked_params)
+          stacked_state = nnx.State.merge(stacked_params, stacked_other)
+        nnx.update(self.local_layers, stacked_state)
+    elif self.local_layers is not None:


Is this the same as line 564?

aireenmei force-pushed the aireen/pr-gemma4-scan-rewrite branch 3 times, most recently from cc77a05 to 2805c03 Compare July 21, 2026 06:09

Rewrite Gemma4 scannable block

ef9cc43

aireenmei force-pushed the aireen/pr-gemma4-scan-rewrite branch from 2805c03 to ef9cc43 Compare July 21, 2026 06:38

aireenmei marked this pull request as ready for review July 21, 2026 06:38

aireenmei assigned NuojCheng Jul 21, 2026

aireenmei assigned RissyRan Jul 21, 2026

aireenmei added the pull ready label Jul 21, 2026

RissyRan added the gemini-review label Jul 22, 2026

RissyRan reviewed Jul 23, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Rewrite Gemma4 scannable block#4530

Rewrite Gemma4 scannable block#4530
aireenmei wants to merge 1 commit into
mainfrom
aireen/pr-gemma4-scan-rewrite

aireenmei commented Jul 18, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Jul 18, 2026

Uh oh!

codecov Bot commented Jul 21, 2026 •

edited

Loading

Uh oh!

aireenmei commented Jul 21, 2026

Uh oh!

github-actions Bot commented Jul 22, 2026

Uh oh!

github-actions Bot commented Jul 22, 2026

Uh oh!

RissyRan left a comment

Uh oh!

RissyRan Jul 23, 2026

Uh oh!

RissyRan Jul 23, 2026

Uh oh!

RissyRan Jul 23, 2026

Uh oh!

RissyRan Jul 23, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

aireenmei commented Jul 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Tests

Checklist

Uh oh!

github-actions Bot commented Jul 18, 2026

🤖 CI Failure Investigation Report

🔍 What Failed

🪵 Error Details & Stack Trace

💡 Root Cause Analysis & Context

🛠️ Recommended Fix

Uh oh!

codecov Bot commented Jul 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

aireenmei commented Jul 21, 2026

Uh oh!

github-actions Bot commented Jul 22, 2026

Uh oh!

github-actions Bot commented Jul 22, 2026

Uh oh!

RissyRan left a comment

Choose a reason for hiding this comment

Uh oh!

RissyRan Jul 23, 2026

Choose a reason for hiding this comment

Uh oh!

RissyRan Jul 23, 2026

Choose a reason for hiding this comment

Uh oh!

RissyRan Jul 23, 2026

Choose a reason for hiding this comment

Uh oh!

RissyRan Jul 23, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

aireenmei commented Jul 18, 2026 •

edited

Loading

codecov Bot commented Jul 21, 2026 •

edited

Loading